import torch
from torch.utils.tensorboard import SummaryWriter

class CNN_NET(torch.nn.Module):
    def __init__(self):
        super(CNN_NET, self).__init__()
        # 定义卷积层
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=3, out_channels=20, kernel_size=5, stride=1, padding=2),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2)
        )
        # 定义全连接层
        self.fc = torch.nn.Sequential(
            torch.nn.Flatten(),
            torch.nn.Dropout(p=0.2),
            torch.nn.Linear(in_features=5000, out_features=2),
            torch.nn.Tanh()
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.fc(x)
        return x

if __name__ == '__main__':
    net = CNN_NET()
    input_img = torch.ones((1,3, 20, 50))
    writer = SummaryWriter('./logs_train')
    writer.add_graph(net, input_img)
    output = net(input_img)
    print(output.shape)
